#include <lcthw/dbg.h>
#include <lcthw/bstree.h>
#include <stdlib.h>
#include <lcthw/bstrlib.h>

// set a function pointer to replace the bstrcmp function
static int default_compare(void *a, void *b)
{
    return bstrcmp((bstring)a, (bstring)b);     // compare two strings deinfed in bstrlib.h
}

// Create a new BSTree
BSTree *BSTree_create(BSTree_compare compare)
{
    BSTree *map = calloc(1, sizeof(BSTree));
    check_mem(map);

    map->compare = compare == NULL ? default_compare : compare;
    return map;

error:
    if (map)
    {
        BSTree_destroy(map);
    }
    return NULL;
}

// Destroy a BSTree node
static int BSTree_destroy_cb(BSTreeNode *node)
{
    free(node);
    return 0;
}

// Destroy a BSTree
void BSTree_destroy(BSTree *map)
{
    if (map)
    {
        BSTree_traverse(map, BSTree_destroy_cb); // Traverse the tree and free each node
        free(map);
    }
}

// set a child node under a designed parent node, the child node don't care whether it is left or right child
static inline BSTreeNode *BSTreeNode_create(BSTreeNode *parent, void *key, void *data) //
{
    BSTreeNode *node = calloc(1, sizeof(BSTreeNode));
    check_mem(node);

    node->key = key;
    node->data = data;
    node->parent = parent;
    return node;

error:
    return NULL;
}

// set a child node to the left or right of the parent node
static inline void BSTree_setnode(BSTree *map, BSTreeNode *node, void *key, void *data)
{
    int cmp = map->compare(node->key, key);

    if (cmp <= 0)
    {
        // if node->left is not NULL, then this node has at least one child node on the left side
        if (node->left)
        {
            BSTree_setnode(map, node->left, key, data); // recursively set the new node to the left of the current node
        }
        else // if node->left is NULL, then this node has no child node
        {
            node->left = BSTreeNode_create(node, key, data); // create a new node and set it to the left of the current node
        }
    }
    else
    {
        if (node->right)
        {
            BSTree_setnode(map, node->right, key, data);
        }
        else
        {
            node->right = BSTreeNode_create(node, key, data);
        }
    }
}

// Set a key-value pair in the BSTree
int BSTree_set(BSTree *map, void *key, void *data)
{
    if (map->root == NULL)
    {
        // if the root node is NULL, then the tree is empty
        map->root = BSTreeNode_create(NULL, key, data); // create a new node and set it as the root node
        check_mem(map->root);
    }
    else
    {
        BSTree_setnode(map, map->root, key, data); // set a new node to the left or right of the root node
    }

    return 0;

error:
    return -1;
}

// get the address of the node with the key in the tree
// if the key is not in the tree, return NULL
static inline BSTreeNode *BSTree_getnode(BSTree *map, BSTreeNode *node, void *key)
{
    int cmp = map->compare(node->key, key);

    if (cmp == 0)
    {
        return node;
    }
    else if (cmp < 0)
    {
        if (node->left)
        {
            return BSTree_getnode(map, node->left, key);
        }
        else
        {
            return NULL;
        }
    }
    else
    {
        if (node->right)
        {
            return BSTree_getnode(map, node->right, key);
        }
        else
        {
            return NULL;
        }
    }
}

// Get the data from a targeted node specified by the key in the BSTree
// if the key is not in the tree, return NULL
void *BSTree_get(BSTree *map, void *key)
{
    if (map->root == NULL)
    {
        return NULL;
    }
    else
    {
        BSTreeNode *node = BSTree_getnode(map, map->root, key);
        return node == NULL ? NULL : node->data;
    }
}


static inline int BSTree_traverse_nodes(BSTreeNode *node, BSTree_traverse_cb traverse_cb)   // traverse the tree
{
    int rc = 0;
    if (node->left)
    {
        rc = BSTree_traverse_nodes(node->left, traverse_cb);
        if (rc != 0)
        {
            return rc;
        }
    }
    if (node->right)
    {
        rc = BSTree_traverse_nodes(node->right, traverse_cb);
        if(rc != 0)
        {
            return rc;
        }
    }
    return traverse_cb(node);
}

// Traverse the tree and apply the callback function to each node
int BSTree_traverse(BSTree *map, BSTree_traverse_cb traverse_cb)
{
    if (map->root)
    {
        return BSTree_traverse_nodes(map->root, traverse_cb);
    }
    return 0;
}

static inline BSTreeNode *BSTree_find_min(BSTreeNode *node)
{
    while (node->left)
    {
        node = node->left;
    }
    return node;
}

// 
static inline void BSTree_replace_node_in_parent(BSTree *map, BSTreeNode *node, BSTreeNode *new_value)
{
    if (node->parent)                                   // if the node has a parent node, it means that node is not the root node
    {
        if (node == node->parent->left)                 // if the node is the left child of the parent node
        {
            node->parent->left = new_value;             // set the new node to the left of the parent node
        }
        else
        {
            node->parent->right = new_value;            // set the new node to the right of the parent node
        }
    }
    else                                                // if the node is the root node
    {
        map->root = new_value;                          // set the new node as the root node                     
    }
    if (new_value)
    {
        new_value->parent = node->parent;               // set the parent node of the new node
    }
}

static inline void BSTree_swap(BSTreeNode *a, BSTreeNode *b)
{
    void *temp = a->key;
    a->key = b->key;
    b->key = temp;

    temp = a->data;
    a->data = b->data;
    b->data = temp;
}

static inline BSTreeNode *BSTree_node_delete(BSTree *map, BSTreeNode *node, void *key)
{
    int cmp = map->compare(node->key, key);

    if (cmp < 0)
    {
        if (node->left)
        {
            return BSTree_node_delete(map, node->left, key);
        }
        else
        {
            return NULL;
        }
    }
    else if (cmp > 0)
    {
        if (node->right)
        {
            return BSTree_node_delete(map, node->right, key);
        }
        else
        {
            return NULL;
        }
    }
    else                                                                        // if the key is found
    {
        if (node->left && node->right)                                          // if the node has two child nodes
        {
            BSTreeNode *successor = BSTree_find_min(node->right);               // find the minimum node in the right subtree
            BSTree_swap(successor, node);
            BSTree_replace_node_in_parent(map, successor, successor->right);    // replace the successor node with its right child node
            return successor;
        }
        else if (node->left)
        {
            BSTree_replace_node_in_parent(map, node, node->left);
        }
        else if (node->right)
        {
            BSTree_replace_node_in_parent(map, node, node->right);
        }
        else
        {
            BSTree_replace_node_in_parent(map, node, NULL);
        }
        return node;
    }
}